from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage,SystemMessage,ToolMessage,AIMessage,BaseMessage
from langchain_core.tools import tool
from langchain_experimental.utilities import PythonREPL
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder

from langgraph.prebuilt import ToolNode
from langgraph.graph import END,StateGraph

from sqlalchemy import create_engine,Column,String,Integer,Float,ForeignKey
from sqlalchemy.orm import sessionmaker,declarative_base
from faker import Faker
from typing import Union,Optional,Annotated,Sequence
from pydantic import BaseModel
from typing_extensions import TypedDict
from PIL import Image as PILImage

import random
import functools
import operator
import io




api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

db_llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base,temperature=0)

coder_llm  = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base,temperature=0)

# response = coder_llm.invoke([HumanMessage(content="帮我写一个使用Python实现的贪吃蛇的游戏代码")])
# print(response)


# 创建基类
Base = declarative_base()

# 定义模型
class SalesData(Base):
    __tablename__ = "sales_data"
    sales_id = Column(Integer,primary_key=True)
    product_id = Column(Integer,ForeignKey("product_information.product_id"))
    employee_id = Column(Integer)
    customer_id = Column(Integer,ForeignKey("customer_information.customer_id"))
    sale_date = Column(String(50))
    quantity = Column(Integer)
    amount = Column(Float)
    discount = Column(Float)

class CustomerInformation(Base):
    __tablename__ = "customer_information"
    customer_id = Column(Integer,primary_key=True)
    customer_name = Column(String(50))
    contact_info = Column(String(50))
    region = Column(String(50))
    customer_type = Column(String(50))

class ProductInformation(Base):
    __tablename__ = "product_information"
    product_id = Column(Integer,primary_key=True)
    product_name = Column(String(50))
    category = Column(String(50))
    unit_price = Column(Float)
    stock_level = Column(Integer)

class CompetitorAnalysis(Base):
    __tablename__ = "competitor_analysis"
    competitor_id = Column(Integer,primary_key=True)
    competitor_name = Column(String(50))
    region = Column(String(50))
    market_share = Column(Float)

# 数据库连接和表创建
#DATABASE_URI = "mysql+pymysql://root:1qaz2wsx@localhost:3306/langgraph?charset=utf8mb4"
DATABASE_URI = "mysql+pymysql://root:root@localhost:3306/langgraph?charset=utf8mb4"
engine = create_engine(DATABASE_URI)
Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)
session = Session()

# 插入模拟数据
# faker = Faker()
# # 生成客户信息
# for _ in range(50):# 生成50个客户
#     customer = CustomerInformation(
#         customer_name=faker.name(),
#         contact_info=faker.phone_number(),
#         region=faker.state(), # 地区
#         customer_type=random.choices(['Retail', 'Wholesale'])  # 零售、批发
#     )
#     session.add(customer)
#
# # 生成产品信息
# for _ in range(20):# 生成20种产品
#     product = ProductInformation(
#         product_name=faker.word(),
#         category=random.choice(['Electronics', 'Clothing', 'Furniture', 'Food', 'Toys']),  # 电子设备，衣服，家具，食品，玩具
#         unit_price=random.uniform(10.0,1000.0),
#         stock_level=random.randint(10,100) # 库存
#     )
#     session.add(product)
#
# # 生成竞争对手信息
# for _ in range(10):# 生成10个竞争对手
#     competitor = CompetitorAnalysis(
#         competitor_name=faker.company(),
#         region=faker.state(),
#         market_share=random.uniform(0.01,0.2) # 市场占有率
#     )
#     session.add(competitor)
#
# # 提交事务
# session.commit()
#
# # 生成销售数据，假设有100条销售记录
# for _ in range(100):
#     sale = SalesData(
#         product_id=random.randint(1,20),
#         employee_id=random.randint(1,20),
#         customer_id=random.randint(1,50),
#         sale_date=faker.date_between(start_date="-1y",end_date="today").strftime("%Y-%m-%d"),
#         quantity=random.randint(1,10),
#         amount=random.uniform(50.0,5000.0),
#         discount=random.uniform(0.0,0.15)
#     )
#     session.add(sale)
#
# session.commit()
# session.close()


class AddSaleSchema(BaseModel):
    product_id:int
    employee_id:int
    customer_id:int
    sale_date:str
    quantity:int
    amount:float
    discount:float

class DeleteSaleSchema(BaseModel):
    sales_id:int

class UpdateSaleSchema(BaseModel):
    sales_id:int
    quantity:int
    amount:float

class QuerySalesSchema(BaseModel):
    sales_id:int

# 1. 添加销售数据
@tool(args_schema=AddSaleSchema)
def add_sale(product_id,employee_id,customer_id,sale_date,quantity,amount,discount):
    """Add sale record to the database."""
    session = Session()
    try:
       new_sale = SalesData(
           product_id = product_id,
           employee_id = employee_id,
           customer_id=customer_id,
           sale_date=sale_date,
           quantity=quantity,
           amount=amount,
           discount=discount
       )
       session.add(new_sale)
       session.commit()
       return {"messages":["销售记录添加成功"]}
    except Exception as e:
        return {"messages":[f"添加失败，错误原因:{e}"]}
    finally:
        session.close()

# 2. 删除销售数据
@tool(args_schema=DeleteSaleSchema)
def delete_sale(sales_id):
    """Delete sale record from the database."""
    session = Session()
    try:
        sale_to_delete = session.query(SalesData).filter(SalesData.sales_id == sales_id).first()
        if sale_to_delete:
            session.delete(sale_to_delete)
            session.commit()
            return {"messages":["销售记录删除成功"]}
        else:
            return {"messages":[f"未找到销售记录ID:{sales_id}"]}
    except Exception as e:
        return {"messages":[f"删除销售数据失败，错误原因:{e}"]}
    finally:
        session.close()

# 3. 修改销售数据
@tool(args_schema=UpdateSaleSchema)
def update_sale(sales_id,quantity,amount):
    """Update sale record in the database."""
    session = Session()
    try:
        sale_to_update = session.query(SalesData).filter(SalesData.sales_id == sales_id).first()
        if sale_to_update:
            sale_to_update.quantity = quantity
            sale_to_update.amount = amount
            session.commit()
    except Exception as e:
        return {"messages":[f"更新销售数据失败,错误原因:{e}"]}
    finally:
        session.close()

# 4. 查询销售数据
@tool(args_schema=QuerySalesSchema)
def query_sales(sales_id):
    """Query sale record from the database."""
    session = Session()
    try:
        sale_data = session.query(SalesData).filter(SalesData.sales_id == sales_id).first()
        if sale_data:
            return {
                "sales_id":sale_data.sales_id,
                "product_id":sale_data.product_id,
                "employee_id":sale_data.employee_id,
                "customer_id":sale_data.customer_id,
                "sale_date":sale_data.sale_date,
                "quantity":sale_data.quantity,
                "amount":sale_data.amount,
                "discount":sale_data.discount
            }
        else:
            return {"messages":[f"未找到销售记录ID:{sales_id}"]}
    except Exception as e:
        return {"messages":[f"查询销售数据失败,错误原因:{e}"]}
    finally:
        session.close()

repl = PythonREPL()

@tool
def python_repl(code:Annotated[str,"The python code to execute to generate your chart."]):
    """Use this to execute python code. If you want to see the output of a value,
        you should print it out with `print(...)`. This is visible to the user."""
    try:
        result = repl.run(code)
    except BaseException as e :
        return f"执行失败,错误原因{e}"
    result_str = f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}"
    return (
            result_str + "\n\nIf you have completed all tasks, respond with FINAL ANSWER."
    )

# 定义工具列表
tools = [add_sale,delete_sale,update_sale,query_sales,python_repl]
tool_executor = ToolNode(tools)

def create_agent(llm,tools,system_message:str):
    """Create an agent."""
    prompt = ChatPromptTemplate(
        [
            (
                "system",
                "You are a helpful AI assistant, collaborating with other assistants."
                " Use the provided tools to progress towards answering the question."
                " If you are unable to fully answer, that's OK, another assistant with different tools "
                " will help where you left off. Execute what you can to make progress."
                " If you or any of the other assistants have the final answer or deliverable,"
                " prefix your response with FINAL ANSWER so the team knows to stop."
                " You have access to the following tools: {tool_names}.\n{system_message}",
            ),
            MessagesPlaceholder(variable_name="messages")
        ]
    )
    prompt = prompt.partial(system_message=system_message)
    prompt = prompt.partial(tool_names=",".join([tool.name for tool in tools]))
    return prompt | llm.bind_tools(tools)

# 数据库管理员
db_agent = create_agent(db_llm,
                        [add_sale,delete_sale,update_sale,query_sales],
                        system_message="You should provide accurate data for the code_generator to use.  and source code shouldn't be the final answer")

# 数据分析师
code_agent = create_agent(coder_llm,[python_repl],system_message="Run python code to display diagrams or output execution results",)

def agent_node(state,agent,name):
    ## 将父图的状态转换为子图的状态，执行出来的是子图的结果
    result = agent.invoke(state)
    # 将代理输出转换为适合附加到全局状态的格式
    if isinstance(result,ToolMessage):
        pass
    else:
        # 创建一个 AIMessage 类的新实例，其中包含 result 对象的所有数据（除了 type 和 name），并且设置新实例的 name 属性为特定的值 name。
        result = AIMessage(**result.model_dump(exclude={"type","name"}),name=name)
    return {"messages":[result],
            # 跟踪发件人，这样我们就知道下一个要传给谁。
            "sender":name
            }

db_node = functools.partial(agent_node,agent=db_agent,name="db_manager")
code_node = functools.partial(agent_node,agent=code_agent,name="code_generator")

# 任何一个代理都可以决定结束
def router(state):
    # 这是一个路由
    messages = state["messages"]
    last_message = messages[-1]
    if hasattr(last_message,"call_tools"):
        # 前一个代理正在调用一个工具
        return "call_tool"
    if "FINAL ANSWER" in last_message.content:
        # 任何Agent都决定工作完成
        return END
    return "continue"

class AgentState(TypedDict):
    messages:Annotated[Sequence[BaseMessage],operator.add]
    sender:str

# 初始化一个状态图
workflow = StateGraph(AgentState)

# 将Agent作为节点进行添加
workflow.add_node("db_manager",db_node)
workflow.add_node("code_generator",code_node)
workflow.add_node("call_tool",tool_executor)

# 通过条件边 构建 子代理之间的通信
workflow.add_conditional_edges(
    "db_manager",
    router,
    {"continue":"code_generator","call_tool":"call_tool",END:END}
)

workflow.add_conditional_edges(
    "code_generator",
    router,
    {"continue":"db_manager","call_tool":"call_tool",END:END}
)

workflow.add_conditional_edges(
    "call_tool",
    lambda x:x["sender"],
    {"db_manager":"db_manager","code_generator":"code_generator"}
)

# 设置 db_manager 为初始节点
workflow.set_entry_point("db_manager")

# 编译图
graph = workflow.compile()

def display(image_data):
    image = PILImage.open(io.BytesIO(image_data))
    image.show()

# display(graph.get_graph().draw_mermaid_png())

for chunk in graph.stream(
                    {"messages":[HumanMessage(content="根据sales_id使用折线图显示前5名销售的销售总额")]},
                        {"recursion_limit":50},
                    stream_mode="values"
                ):
    print(chunk)